実験管理が簡単に行えるmlflow trackingをローカル環境上で試してみた。

2018年にDatabricks社が開発し、OSSのライブラリとして提供しています。
現在も活発にコミュニティによってメンテナンスされています。

Mlflowライブラリは、機械学習を実行するものではないことに注意して下さい。
基本的に機械学習ライブラリ(sklearn, pytorch, tensolflow等)と併用して使うものです。

[[source]]
url = "https://pypi.org/simple"
verify_ssl = true
name = "pypi"

[packages]
pandas = "==2.0.3"
numpy = "==1.24.3"
scikit-learn = "==1.3.0"
lightgbm = "==4.0.0"
mlflow = "==2.8.0"
optuna = "==3.6.1"

[dev-packages]

[requires]
python_version = "3.11"
python_full_version = "3.11.9"

├── titanic-mlflow
│   ├── data # 学習/テスト/提出用データ等の保管フォルダ
│   │   ├── gender_submission.csv
│   │   ├── test.csv
│   │   └── train.csv
│   ├── plots # 学習過程や結果を可視化する画像フォルダ
│   └── src   # 前処理からモデル構築を実装する.pyフォルダ
│       └── main.py
├── .python-version
├── Pipfile
└── Pipfile.lock

import os
import warnings
import sys
from datetime import datetime

import pandas as pd
import numpy as np
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
from sklearn.model_selection import train_test_split
import lightgbm as lgb
from urllib.parse import urlparse
import mlflow
import mlflow.lightgbm
from mlflow.exceptions import MlflowException

from mlflow.models import ModelSignature
from mlflow.types.schema import Schema, ColSpec
from mlflow.types import DataType

import optuna
import matplotlib.pyplot as plt

def eval_metrics(actual, pred):
    """評価指標を計算する関数。

    Args:
        actual (array-like): 実際の値。
        pred (array-like): 予測値。

    Returns:
        tuple: 正確度、適合率、再現率、F1スコアのタプル。
    """
    accuracy = accuracy_score(actual, pred)
    precision = precision_score(actual, pred)
    recall = recall_score(actual, pred)
    f1 = f1_score(actual, pred)
    return accuracy, precision, recall, f1

def load_and_preprocess_data(csv_path):
    """データを読み込み、前処理を行う関数。

    Args:
        csv_path (str): CSVファイルのパス。

    Returns:
        tuple: 前処理された訓練データとテストデータ。
    """
    print(f"Loading data from: {csv_path}")

    if not os.path.exists(csv_path):
        raise FileNotFoundError(f"ファイルが見つかりません: {csv_path}")

    data = pd.read_csv(csv_path)
    data = data.drop(['Name', 'Ticket', 'Cabin'], axis=1)
    data = data.dropna()
    data['Sex'] = data['Sex'].map({'male': 0, 'female': 1})
    data['Embarked'] = data['Embarked'].map({'C': 0, 'Q': 1, 'S': 2})

    train, test = train_test_split(data, test_size=0.2, random_state=42)

    train_x = train.drop(["Survived"], axis=1)
    test_x = test.drop(["Survived"], axis=1)
    train_y = train["Survived"]
    test_y = test["Survived"]

    return train_x, test_x, train_y, test_y

def log_mlflow(trial, train_x, train_y, test_x, test_y, base_dir):
    """MLflowにパラメータとメトリクスをログする関数。

    Args:
        trial (optuna.trial.FrozenTrial): Optunaのトライアルオブジェクト。
        train_x (DataFrame): 訓練データの特徴量。
        train_y (Series): 訓練データのターゲット。
        test_x (DataFrame): テストデータの特徴量。
        test_y (Series): テストデータのターゲット。
        base_dir (str): ベースディレクトリのパス。
    """
    # アクティブなランがある場合は終了する
    if mlflow.active_run() is not None:
        mlflow.end_run()

    run_name = f"Trial {trial.number + 1} - {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}"
    with mlflow.start_run(run_name=run_name) as run:
        lgbm = lgb.LGBMClassifier(
            num_leaves=trial.params["num_leaves"],
            learning_rate=trial.params["learning_rate"],
            n_estimators=trial.params["n_estimators"],
            random_state=42,
            force_row_wise=True
        )

        # 学習曲線のためにevals_resultを使用
        evals_result = {}
        lgbm.fit(
            train_x, train_y,
            eval_set=[(train_x, train_y), (test_x, test_y)],
            eval_metric='logloss',
            callbacks=[lgb.record_evaluation(evals_result)]
        )

        predicted_survival = lgbm.predict(test_x)

        (accuracy, precision, recall, f1) = eval_metrics(test_y, predicted_survival)

        print(f"LightGBM model (num_leaves={trial.params['num_leaves']}, learning_rate={trial.params['learning_rate']}, n_estimators={trial.params['n_estimators']}):")
        print("  Accuracy: %s" % accuracy)
        print("  Precision: %s" % precision)
        print("  Recall: %s" % recall)
        print("  F1 Score: %s" % f1)

        # 学習パラメータと精度メトリックスをTracking
        mlflow.log_param("num_leaves", trial.params["num_leaves"])
        mlflow.log_param("learning_rate", trial.params["learning_rate"])
        mlflow.log_param("n_estimators", trial.params["n_estimators"])
        mlflow.log_metric("accuracy", accuracy)
        mlflow.log_metric("precision", precision)
        mlflow.log_metric("recall", recall)
        mlflow.log_metric("f1", f1)

        mlflow.set_tag("release.version", "1.0.0")

        # アーティファクトのログ
        mlflow.log_artifact(os.path.join(base_dir, "../data/train.csv"), "data")
        mlflow.log_artifact(os.path.join(base_dir, "../data/test.csv"), "data")
        mlflow.log_artifact(os.path.join(base_dir, "../data/gender_submission.csv"), "data")
        mlflow.log_artifact(os.path.join(base_dir, "../src/main.py"), "src")

        # モデルスキーマの定義
        input_schema = Schema([
            ColSpec(DataType.integer, "Pclass"),
            ColSpec(DataType.float, "Age"),
            ColSpec(DataType.integer, "SibSp"),
            ColSpec(DataType.integer, "Parch"),
            ColSpec(DataType.float, "Fare"),
            ColSpec(DataType.integer, "Sex"),
            ColSpec(DataType.integer, "Embarked")
        ])
        output_schema = Schema([ColSpec(DataType.integer, "Survived")])
        signature = ModelSignature(inputs=input_schema, outputs=output_schema)

        # モデルの保存
        mlflow.lightgbm.log_model(lgbm, "model", signature=signature)
        print("Model saved in run %s" % mlflow.active_run().info.run_uuid)

        # 学習曲線のプロット
        plt.figure(figsize=(10, 5))
        plt.plot(evals_result['training']['binary_logloss'], label='train')
        plt.plot(evals_result['valid_1']['binary_logloss'], label='test')
        plt.xlabel('Iteration')
        plt.ylabel('Log Loss')
        plt.title('Learning Curve')
        plt.legend()
        plt.grid(True)

        # 学習曲線の保存
        plots_dir = os.path.join(base_dir, "../plots")
        os.makedirs(plots_dir, exist_ok=True)
        learning_curve_path = os.path.join(plots_dir, f"learning_curve_trial_{trial.number + 1}.png")
        plt.savefig(learning_curve_path)
        mlflow.log_artifact(learning_curve_path, "plots")
        plt.close()

    # ランを終了
    mlflow.end_run()

def objective(trial, train_x, train_y, test_x, test_y):
    """Optunaの目的関数。

    Args:
        trial (optuna.trial.Trial): Optunaのトライアルオブジェクト。
        train_x (DataFrame): 訓練データの特徴量。
        train_y (Series): 訓練データのターゲット。
        test_x (DataFrame): テストデータの特徴量。
        test_y (Series): テストデータのターゲット。

    Returns:
        float: 負の正確度。
    """
    num_leaves = trial.suggest_int("num_leaves", 20, 50)
    learning_rate = trial.suggest_float("learning_rate", 0.01, 0.1)
    n_estimators = trial.suggest_int("n_estimators", 50, 200)

    lgbm = lgb.LGBMClassifier(
        num_leaves=num_leaves,
        learning_rate=learning_rate,
        n_estimators=n_estimators,
        random_state=42,
        force_row_wise=True
    )

    lgbm.fit(train_x, train_y)
    predicted_survival = lgbm.predict(test_x)
    accuracy, precision, recall, f1 = eval_metrics(test_y, predicted_survival)

    # Optunaは目的関数が最小化されるように設計されているため、負の精度を返す
    return -accuracy

def main():
    """メイン関数。データの前処理、Optunaによるハイパーパラメータ最適化、MLflowによるログ記録を行う。
    """
    warnings.filterwarnings("ignore")
    np.random.seed(40)

    # トラッキングサーバの場所
    tracking_uri = mlflow.get_tracking_uri()
    print('Current tracking uri: {}'.format(tracking_uri))
    # レジストリサーバの場所
    mr_uri = mlflow.get_registry_uri()
    print('Current model registry uri: {}'.format(mr_uri))
    # Artifactストレージの場所
    artifact_uri = mlflow.get_artifact_uri()
    print('Current artifact uri: {}'.format(artifact_uri))

    # スクリプトのディレクトリを基準にパスを解決
    base_dir = os.path.abspath(os.path.dirname(__file__))
    csv_path = os.path.join(base_dir, "../data/train.csv")  # Titanicデータのパスを指定

    # デバッグのためにファイルパスを出力
    print(f"Loading data from: {csv_path}")

    train_x, test_x, train_y, test_y = load_and_preprocess_data(csv_path)

    experiment_name = "titanic_classification"
    try:
        # 指定された名前の新しい実験を作成
        mlflow.create_experiment(experiment_name)
        mlflow.set_experiment(experiment_name)
        print(f"Experiment '{experiment_name}' created.")
    except MlflowException as e:
        # 実験が既に存在する場合はその実験を設定
        if "already exists" in str(e):
            mlflow.set_experiment(experiment_name)
            print(f"Experiment '{experiment_name}' already exists. Set the existing experiment.")
        else:
            # その他の MlflowException が発生した場合は再スロー
            raise

    study = optuna.create_study(direction="maximize")
    study.optimize(lambda trial: objective(trial, train_x, train_y, test_x, test_y), n_trials=5)

    best_trial = study.best_trial
    print(f"Best trial: {best_trial.params}")

    for trial in study.trials:
        log_mlflow(trial, train_x, train_y, test_x, test_y, base_dir)

if __name__ == "__main__":
    main()

概要
評価指標を計算する関数。

1. 正確度 (Accuracy)
2. 適合率 (Precision)
3. 再現率 (Recall)
4. F1スコア (F1 Score)

概要
データを読み込み、前処理を行う関数。

1. データの読み込み。
2. 不要な列の削除 (Name, Ticket, Cabin)。
3. 欠損値の削除。
4. カテゴリカル変数 (Sex, Embarked) の数値変換。
5. 訓練データとテストデータに分割。

概要
Optunaの目的関数。

1. ハイパーパラメータのサンプリング。
2. LightGBMモデルの学習と評価。
3. 評価指標（正確度）の計算。
4. 負の正確度を返す（Optunaは目的関数が最小化されるように設計されているため）。

概要
MLflowにパラメータとメトリクスをログする関数。

1. アクティブなランがある場合は終了。
2. 新しいランを開始し、モデルの学習と評価を実行。
3. 学習結果と評価指標をMLflowにログ。
4. 学習曲線のプロットを作成し、保存。
5. モデルとそのスキーマをMLflowに保存。

# アクティブなランがある場合は終了する
if mlflow.active_run() is not None:
    mlflow.end_run()

run_name = f"Trial {trial.number + 1} - {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}"
with mlflow.start_run(run_name=run_name) as run:
    lgbm = lgb.LGBMClassifier(
        num_leaves=trial.params["num_leaves"],
        learning_rate=trial.params["learning_rate"],
        n_estimators=trial.params["n_estimators"],
        random_state=42,
        force_row_wise=True
    )

    # 学習曲線のためにevals_resultを使用
    evals_result = {}
    lgbm.fit(
        train_x, train_y,
        eval_set=[(train_x, train_y), (test_x, test_y)],
        eval_metric='logloss',
        callbacks=[lgb.record_evaluation(evals_result)]
    )

    predicted_survival = lgbm.predict(test_x)

    (accuracy, precision, recall, f1) = eval_metrics(test_y, predicted_survival)

    print(f"LightGBM model (num_leaves={trial.params['num_leaves']}, learning_rate={trial.params['learning_rate']}, n_estimators={trial.params['n_estimators']}):")
    print("  Accuracy: %s" % accuracy)
    print("  Precision: %s" % precision)
    print("  Recall: %s" % recall)
    print("  F1 Score: %s" % f1)

# 学習パラメータと精度メトリックスをTracking
mlflow.log_param("num_leaves", trial.params["num_leaves"])
mlflow.log_param("learning_rate", trial.params["learning_rate"])
mlflow.log_param("n_estimators", trial.params["n_estimators"])
mlflow.log_metric("accuracy", accuracy)
mlflow.log_metric("precision", precision)
mlflow.log_metric("recall", recall)
mlflow.log_metric("f1", f1)

mlflow.set_tag("release.version", "1.0.0")

# アーティファクトのログ
mlflow.log_artifact(os.path.join(base_dir, "../data/train.csv"), "data")
mlflow.log_artifact(os.path.join(base_dir, "../data/test.csv"), "data")
mlflow.log_artifact(os.path.join(base_dir, "../data/gender_submission.csv"), "data")
mlflow.log_artifact(os.path.join(base_dir, "../src/main.py"), "src")

# 学習曲線のプロット
plt.figure(figsize=(10, 5))
plt.plot(evals_result['training']['binary_logloss'], label='train')
plt.plot(evals_result['valid_1']['binary_logloss'], label='test')
plt.xlabel('Iteration')
plt.ylabel('Log Loss')
plt.title('Learning Curve')
plt.legend()
plt.grid(True)

# 学習曲線の保存
plots_dir = os.path.join(base_dir, "../plots")
os.makedirs(plots_dir, exist_ok=True)
learning_curve_path = os.path.join(plots_dir, f"learning_curve_trial_{trial.number + 1}.png")
plt.savefig(learning_curve_path)
mlflow.log_artifact(learning_curve_path, "plots")
plt.close()

# モデルスキーマの定義
input_schema = Schema([
    ColSpec(DataType.integer, "Pclass"),
    ColSpec(DataType.float, "Age"),
    ColSpec(DataType.integer, "SibSp"),
    ColSpec(DataType.integer, "Parch"),
    ColSpec(DataType.float, "Fare"),
    ColSpec(DataType.integer, "Sex"),
    ColSpec(DataType.integer, "Embarked")
])
output_schema = Schema([ColSpec(DataType.integer, "Survived")])
signature = ModelSignature(inputs=input_schema, outputs=output_schema)

# モデルの保存
mlflow.lightgbm.log_model(lgbm, "model", signature=signature)
print("Model saved in run %s" % mlflow.active_run().info.run_uuid)

概要
メイン関数。データの前処理、Optunaによるハイパーパラメータ最適化、MLflowによるログ記録を行う。

1. ランダムシードの設定。
2. MLflowのトラッキングURI、モデルレジストリURI、アーティファクトURIの取得と表示。
3. データの読み込みと前処理。
4. MLflow実験の作成または設定。
5. Optunaによるハイパーパラメータ最適化の実行。
6. 各トライアルの結果をMLflowにログ。

np.random.seed(40)

# トラッキングサーバの場所(実験結果が保存されるサーバ)
tracking_uri = mlflow.get_tracking_uri()
print('Current tracking uri: {}'.format(tracking_uri))
# レジストリサーバの場所(モデルが登録されるサーバ)
mr_uri = mlflow.get_registry_uri()
print('Current model registry uri: {}'.format(mr_uri))
# Artifactストレージの場所(トレーニングデータやモデルファイルなどが保存されるストレージ)
artifact_uri = mlflow.get_artifact_uri()
print('Current artifact uri: {}'.format(artifact_uri))

# スクリプトのディレクトリを基準にパスを解決
base_dir = os.path.abspath(os.path.dirname(__file__))
csv_path = os.path.join(base_dir, "../data/train.csv")  # Titanicデータのパスを指定

# デバッグのためにファイルパスを出力
print(f"Loading data from: {csv_path}")

train_x, test_x, train_y, test_y = load_and_preprocess_data(csv_path)

experiment_name = "titanic_classification"
try:
    # 指定された名前の新しい実験を作成
    mlflow.create_experiment(experiment_name)
    mlflow.set_experiment(experiment_name)
    print(f"Experiment '{experiment_name}' created.")
except MlflowException as e:
    # 実験が既に存在する場合はその実験を設定
    if "already exists" in str(e):
        mlflow.set_experiment(experiment_name)
        print(f"Experiment '{experiment_name}' already exists. Set the existing experiment.")
    else:
        # その他の MlflowException が発生した場合は再スロー
        raise

study = optuna.create_study(direction="maximize")
study.optimize(lambda trial: objective(trial, train_x, train_y, test_x, test_y), n_trials=5)

best_trial = study.best_trial
print(f"Best trial: {best_trial.params}")

for trial in study.trials:
    log_mlflow(trial, train_x, train_y, test_x, test_y, base_dir)

python titanic-mlflow/src/main.py

mlflow ui

pip install mlflow

import mlflow # 追加
import mlflow.sklearn # 追加
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score

# データの読み込み
iris = load_iris()
X_train, X_test, y_train, y_test = train_test_split(iris.data, iris.target, test_size=0.2, random_state=42)

# MLflowのランを開始
with mlflow.start_run(): # 追加
    # モデルのトレーニング
    clf = RandomForestClassifier(n_estimators=100, random_state=42)
    clf.fit(X_train, y_train)

    # 予測と評価
    predictions = clf.predict(X_test)
    accuracy = accuracy_score(y_test, predictions)

    # パラメータとメトリクスをログ
    mlflow.log_param("n_estimators", 100) # 追加
    mlflow.log_metric("accuracy", accuracy) # 追加

    # モデルをログ
    mlflow.sklearn.log_model(clf, "model") # 追加

    print(f"Model accuracy: {accuracy}")